import json
import time
import hashlib
import base64
import hmac
import ssl
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
from wsgiref.handlers import format_date_time
import websocket
import pymysql
import jieba
import jieba.analyse
import requests
import threading
import urllib.parse  # 添加这行导入

class SparkAIClient:
    def __init__(self, app_id, api_secret, api_key, version="ultra"):
        """
        初始化讯飞星火AI客户端 - 使用WebSocket方式
        """
        self.app_id = app_id
        self.api_secret = api_secret
        self.api_key = api_key
        self.version = version

        # 根据版本设置host和path
        if version == "lite":
            self.host = "spark-api.xf-yun.com"
            self.path = "/v1.1/chat"
            self.domain = "lite"
        elif version == "pro":
            self.host = "spark-api.xf-yun.com"
            self.path = "/v3.1/chat"
            self.domain = "generalv3"
        elif version == "max":
            self.host = "spark-api.xf-yun.com"
            self.path = "/v3.5/chat"
            self.domain = "generalv3.5"
        elif version == "ultra":
            self.host = "spark-api.xf-yun.com"
            self.path = "/v4.0/chat"
            self.domain = "4.0Ultra"
        else:
            # 默认使用ultra
            self.host = "spark-api.xf-yun.com"
            self.path = "/v4.0/chat"
            self.domain = "4.0Ultra"

        self.spark_url = f"wss://{self.host}{self.path}"

    def _create_signature(self):
        """修复签名生成方法"""
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 构造签名原始字符串
        signature_origin = f"host: {self.host}\ndate: {date}\nGET {self.path} HTTP/1.1"

        # HMAC-SHA256 加密
        signature_sha = hmac.new(
            self.api_secret.encode('utf-8'),
            signature_origin.encode('utf-8'),
            digestmod=hashlib.sha256
        ).digest()

        signature_sha_base64 = base64.b64encode(signature_sha).decode('utf-8')

        # 构造 authorization
        authorization_origin = f'api_key="{self.api_key}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode('utf-8')

        return authorization, date

    def _build_url(self):
        """构建URL - 修复编码问题"""
        authorization, date = self._create_signature()

        # 正确的参数编码
        params = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }

        # 使用urlencode进行正确编码
        query_string = urlencode(params, safe='')
        url = f"wss://{self.host}{self.path}?{query_string}"

        print(f"最终URL: {url}")
        return url

    def ask(self, prompt, temperature=0.5, max_tokens=2048):
        """改进的ask方法，正确处理流式响应"""
        try:
            answer = ""
            status_completed = False

            def on_message(ws, message):
                nonlocal answer, status_completed
                # print(f"收到消息: {message}")
                try:
                    data = json.loads(message)
                    header = data.get('header', {})
                    code = header.get('code')

                    if code != 0:
                        print(f"API返回错误: {header.get('message', '未知错误')}")
                        return

                    # 提取内容
                    payload = data.get('payload', {})
                    choices = payload.get('choices', {})
                    text_list = choices.get('text', [{}])

                    if text_list:
                        content = text_list[0].get('content', '')
                        if content:
                            answer += content

                    # 检查是否完成
                    status = choices.get('status')
                    if status == 2:
                        status_completed = True
                        ws.close()

                except Exception as e:
                    print(f"解析消息错误: {e}")

            def on_error(ws, error):
                print(f"WebSocket错误: {error}")

            def on_close(ws, close_status_code, close_msg):
                print("WebSocket连接关闭")

            def on_open(ws):
                # 构建正确的请求数据
                request_data = {
                    "header": {
                        "app_id": self.app_id,
                        "uid": "user_" + str(int(time.time()))
                    },
                    "parameter": {
                        "chat": {
                            "domain": self.domain,
                            "temperature": temperature,
                            "max_tokens": max_tokens
                        }
                    },
                    "payload": {
                        "message": {
                            "text": [
                                {
                                    "role": "user",
                                    "content": prompt
                                }
                            ]
                        }
                    }
                }
                print(f"发送请求数据: {json.dumps(request_data, ensure_ascii=False)}")
                ws.send(json.dumps(request_data, ensure_ascii=False))

            ws_url = self._build_url()
            print(f"连接URL: {ws_url}")

            ws = websocket.WebSocketApp(
                ws_url,
                on_message=on_message,
                on_error=on_error,
                on_close=on_close,
                on_open=on_open
            )

            # 设置更长的超时时间
            ws.run_forever(
                sslopt={"cert_reqs": ssl.CERT_NONE},
                ping_interval=30,
                ping_timeout=10
            )

            # 等待响应完成
            start_time = time.time()
            while not status_completed and (time.time() - start_time) < 30:  # 30秒超时
                time.sleep(0.1)

            return answer if answer else "抱歉，暂时无法获取回答"

        except Exception as e:
            print(f"WebSocket通信错误: {str(e)}")
            return f"错误: {str(e)}"


class SimpleAIClient:
    """简单AI客户端 - 用于测试和备用"""

    def __init__(self):
        self.responses = {
            "你好": "你好！我是招聘数据分析助手，可以帮您分析招聘市场趋势、查询职位信息等。",
            "招聘": "我可以帮您分析招聘数据，包括：\n- 各城市职位分布\n- 薪资水平分析\n- 技能要求统计\n- 公司招聘趋势",
            "薪资": "根据当前数据，我可以为您提供：\n- 各岗位薪资范围\n- 不同城市薪资对比\n- 经验要求的薪资差异",
            "职位": "我可以帮您查询：\n- 特定技能的职位需求\n- 各公司的招聘情况\n- 热门技术栈分析"
        }

    def ask(self, prompt):
        """简单回复"""
        prompt_lower = prompt.lower()

        for key in self.responses:
            if key in prompt_lower:
                return self.responses[key]

        # 默认回复
        return f"您好！关于\"{prompt}\"，作为招聘数据分析助手，我可以帮您：\n\n1. 分析各城市招聘职位分布\n2. 查询特定技能的薪资水平\n3. 比较不同公司的招聘要求\n4. 提供就业市场趋势分析\n\n请告诉我您想了解的具体内容。"


class JobKnowledgeBase:
    def __init__(self, db_config):
        self.db_config = db_config

    def get_db_connection(self):
        return pymysql.connect(**self.db_config)

    def extract_keywords(self, question):
        """从问题中提取关键词"""
        try:
            keywords = jieba.analyse.extract_tags(question, topK=5, withWeight=False)
            return keywords
        except:
            # 简单的关键词提取备用方案
            words = question.replace('?', ' ').replace('？', ' ').split()
            return words[:5]

    def search_relevant_jobs(self, keywords):
        """根据关键词搜索相关职位"""
        conn = self.get_db_connection()
        cursor = conn.cursor(pymysql.cursors.DictCursor)

        if not keywords:
            return []

        try:
            # 构建查询条件
            conditions = []
            params = []
            for kw in keywords:
                conditions.append("""
                    job_title LIKE %s OR 
                    job_company LIKE %s OR 
                    job_skill LIKE %s OR 
                    keyword LIKE %s OR 
                    company_city LIKE %s
                """)
                params.extend([f'%{kw}%'] * 5)

            where_clause = " OR ".join(conditions)
            sql = f"""
            SELECT 
                job_title, job_salary, job_company, company_city, job_skill, 
                DATE_FORMAT(crawl_time, '%%Y-%%m-%%d') as crawl_date
            FROM bosszp 
            WHERE {where_clause}
            ORDER BY crawl_time DESC 
            LIMIT 10
            """

            cursor.execute(sql, params)
            jobs = cursor.fetchall()
            return jobs

        except Exception as e:
            print(f"搜索职位失败: {e}")
            return []
        finally:
            cursor.close()
            conn.close()

    def get_statistics(self):
        """获取数据库统计信息"""
        conn = self.get_db_connection()
        cursor = conn.cursor(pymysql.cursors.DictCursor)

        try:
            # 获取总职位数
            cursor.execute("SELECT COUNT(*) as total_jobs FROM bosszp")
            total_jobs = cursor.fetchone()['total_jobs']

            # 获取城市分布
            cursor.execute("""
                SELECT company_city, COUNT(*) as count 
                FROM bosszp 
                GROUP BY company_city 
                ORDER BY count DESC 
                LIMIT 10
            """)
            city_distribution = cursor.fetchall()

            # 获取关键词分布
            cursor.execute("""
                SELECT keyword, COUNT(*) as count 
                FROM bosszp 
                GROUP BY keyword 
                ORDER BY count DESC 
                LIMIT 10
            """)
            keyword_distribution = cursor.fetchall()

            return {
                'total_jobs': total_jobs,
                'city_distribution': city_distribution,
                'keyword_distribution': keyword_distribution
            }

        except Exception as e:
            print(f"获取统计信息失败: {e}")
            return {'total_jobs': 0, 'city_distribution': [], 'keyword_distribution': []}
        finally:
            cursor.close()
            conn.close()

    def is_recruitment_question(self, question):
        """判断问题是否与招聘相关"""
        recruitment_keywords = [
            '职位', '工作', '招聘', '薪资', '工资', '公司', '入职', '求职',
            '面试', '简历', '岗位', '入职', '开发', '工程师', '技术', '经验',
            '技能', '要求', '待遇', '福利', '加班', '双休', '五险一金',
            'java', 'python', '前端', '后端', '算法', '测试', '运维'
        ]

        question_lower = question.lower()
        return any(keyword in question_lower for keyword in recruitment_keywords)

    def build_context(self, question):
        """构建AI助手的上下文"""
        is_recruitment = self.is_recruitment_question(question)

        if is_recruitment:
            # 招聘相关问题：使用数据库信息
            keywords = self.extract_keywords(question)
            jobs = self.search_relevant_jobs(keywords)
            stats = self.get_statistics()

            context = f"""你是一个专业的招聘助手，专门帮助用户解答关于招聘信息的问题。

用户的问题: {question}

"""

            if stats['total_jobs'] > 0:
                context += f"当前数据库中有 {stats['total_jobs']} 个职位信息。"

                if stats['city_distribution']:
                    cities = [item['company_city'] for item in stats['city_distribution'][:5]]
                    context += f"覆盖城市: {', '.join(cities)}。"

                if stats['keyword_distribution']:
                    keywords = [item['keyword'] for item in stats['keyword_distribution'][:5]]
                    context += f"热门关键词: {', '.join(keywords)}。"

            if jobs:
                context += "\n\n以下是相关的招聘信息：\n"
                for i, job in enumerate(jobs):
                    context += f"{i + 1}. 职位: {job['job_title']}, 公司: {job['job_company']}, "
                    context += f"城市: {job['company_city']}, 薪资: {job['job_salary']}\n"

                context += "\n请根据以上信息回答用户的问题。如果信息不足，请基于你的知识进行回答。"
            else:
                context += "\n当前没有找到完全匹配的招聘信息，请基于你的知识进行回答。"
        else:
            # 非招聘问题：直接回答
            context = f"用户的问题: {question}\n请以专业助手的身份回答用户的问题。"

        return context, is_recruitment


class AIAssistant:
    """AI助手主类"""

    def __init__(self, db_config, spark_config=None, use_simple=False):
        self.knowledge_base = JobKnowledgeBase(db_config)
        self.use_simple = use_simple

        if use_simple:
            self.ai_client = SimpleAIClient()
            print("使用简单AI客户端")
        elif spark_config and all(spark_config.values()):
            try:
                self.ai_client = SparkAIClient(
                    app_id=spark_config['app_id'],
                    api_secret=spark_config['api_secret'],
                    api_key=spark_config['api_key'],
                    version="ultra"
                )
                print("讯飞星火客户端初始化成功")
            except Exception as e:
                print(f"讯飞星火客户端初始化失败: {e}, 回退到简单客户端")
                self.ai_client = SimpleAIClient()
                self.use_simple = True
        else:
            self.ai_client = SimpleAIClient()
            self.use_simple = True
            print("使用简单AI客户端")

    def ask(self, question):
        """改进的提问方法"""
        # 构建上下文
        context, is_recruitment = self.knowledge_base.build_context(question)

        print(f"问题: {question}")
        print(f"上下文长度: {len(context)}")
        print(f"问题类型: {'招聘相关' if is_recruitment else '日常问题'}")
        print(f"使用客户端: {'简单客户端' if self.use_simple else '讯飞星火'}")

        try:
            if self.use_simple:
                # 简单客户端直接使用原始问题
                answer = self.ai_client.ask(question)
            else:
                # 星火客户端使用构建的上下文
                print(f"发送给星火的上下文: {context[:500]}...")  # 只打印前500字符
                answer = self.ai_client.ask(context)
                print(f"星火返回的答案: {answer}")

            if not answer or len(answer.strip()) == 0:
                answer = "抱歉，暂时无法回答这个问题"
                # 如果星火API失败，回退到简单客户端
                if not self.use_simple:
                    print("星火API返回空答案，回退到简单客户端")
                    answer = SimpleAIClient().ask(question)

            return {
                'answer': answer,
                'is_recruitment': is_recruitment,
                'ai_client': 'simple' if self.use_simple else 'spark'
            }

        except Exception as e:
            print(f"AI客户端错误: {e}")
            # 出错时回退到简单客户端
            try:
                backup_answer = SimpleAIClient().ask(question)
                return {
                    'answer': backup_answer,
                    'is_recruitment': is_recruitment,
                    'ai_client': 'backup'
                }
            except:
                return {
                    'answer': "抱歉，AI助手暂时无法回答。请稍后再试。",
                    'is_recruitment': is_recruitment,
                    'ai_client': 'error'
                }

def create_ai_assistant(db_config, use_backup=False):
    """创建AI助手实例"""

    # 使用您提供的有效密钥
    spark_config = {
        'app_id': '0adb1398',
        'api_secret': 'MGVkMjE3OTQwZTdlMzIzNDRhNzRmOTNh',
        'api_key': '4e0ab6929faad9ff81b6b84d0fc8bf76'
    }

    # 如果明确要求使用备用模式，或者密钥无效，则使用简单客户端
    if use_backup or not all(spark_config.values()):
        print("使用备用模式（简单AI客户端）")
        return AIAssistant(db_config, use_simple=True)
    else:
        print("尝试使用讯飞星火API")
        return AIAssistant(db_config, spark_config=spark_config, use_simple=False)


# 测试函数
def test_ai_assistant():
    """测试AI助手"""
    db_config = {
        'host': '127.0.0.1',
        'port': 3306,
        'user': 'root',
        'password': '123456',
        'database': 'recruitment_analysis',
        'charset': 'utf8mb4'
    }

    print("=== 测试简单AI客户端 ===")
    assistant_simple = create_ai_assistant(db_config, use_backup=True)

    test_questions = [
        "你好",
        "Java开发的薪资怎么样？",
        "广州有哪些招聘职位？"
    ]

    for question in test_questions:
        print(f"\n问题: {question}")
        result = assistant_simple.ask(question)
        print(f"回答: {result['answer']}")
        print(f"类型: {'招聘相关' if result['is_recruitment'] else '日常问题'}")
        print(f"AI客户端: {result['ai_client']}")

    print("\n=== 测试讯飞星火API ===")
    try:
        assistant_spark = create_ai_assistant(db_config, use_backup=False)
        result = assistant_spark.ask("鲁迅是谁？")
        print(f"星火API测试结果: {result}")
    except Exception as e:
        print(f"星火API测试失败: {e}")


if __name__ == "__main__":
    test_ai_assistant()